package com.hjl.sr.security;

import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;

import com.hjl.sr.remoting.EncryptContext;
import com.hjl.sr.security.token.TokenManager;
import com.hjl.sr.util.BytesUtil;
public class PrefixTokenEncrypter implements Encrypter {

	
	private Encrypter encrypter;
	
	private TokenManager tokenManager;

	public PrefixTokenEncrypter(Encrypter encrypter, TokenManager tokenManager) {
		super();
		this.encrypter = encrypter;
		this.tokenManager = tokenManager;
	}

	public InputStream decryptInputStream(InputStream inputStream, EncryptContext encryptContext) {
		String token = tokenManager.getDecryptToken(encryptContext);
		final byte[] prefixBytes = BytesUtil.getBytes(token) ;
		final int prefixBytesLength;
		if (prefixBytes == null) {
			prefixBytesLength = 0;
		} else {
			prefixBytesLength = prefixBytes.length;
		}
		InputStream is = encrypter.decryptInputStream(inputStream, encryptContext);
		if (prefixBytesLength == 0) {
			return is;
		}
		for (int i = 0; i < prefixBytesLength; ++i) {
			int c;
			try {
				c = is.read();
			} catch (IOException e) {
				throw new RuntimeException(e);
			}
			if (c < 0) {
				throw new SecurityException("data length too less");
			}
			if (c != prefixBytes[i]) {
				throw new SecurityException("data token not match");
			}
		}
		return is;
	}

	public OutputStream encryptOutputStream(OutputStream outputStream, EncryptContext encryptContext) {
		String token = tokenManager.getEncryptToken(encryptContext);
		final byte[] prefixBytes = BytesUtil.getBytes(token);
		try {
			final int prefixBytesLength;
			if (prefixBytes == null) {
				prefixBytesLength = 0;
			} else {
				prefixBytesLength = prefixBytes.length;
			}
			OutputStream os = encrypter.encryptOutputStream(outputStream, encryptContext);
			if (prefixBytesLength == 0) {
				return os;
			}
			try {
				os.write(prefixBytes);
			} catch (IOException e) {
				throw new RuntimeException(e);
			}
			return os;
		} finally {
			tokenManager.clearToken(encryptContext, token);
		}
	}

}
